﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Expectation_Maximization
{
	class ExpectationMaximizer
	{

		private int _motifSize;
		private Dictionary<string, int[]> _residueCounts;
		private Dictionary<string, double[]> _residueFrequencies;
		private Dictionary<string, double[]> _logOddsScores;
		private List<HelperFunctions.Sequence> _sequences;
		private List<int[]> _startingAlignments;
		private double _maxOddsLogScore;
		private List<HelperFunctions.Sequence> _maxOddsLogScoreSequences;

		public ExpectationMaximizer(List<HelperFunctions.Sequence> sequences, int motifSize)
		{
			_residueCounts = new Dictionary<string, int[]>();
			_residueFrequencies = new Dictionary<string, double[]>();
			_logOddsScores = new Dictionary<string, double[]>();
			_motifSize = motifSize;
			_sequences = sequences;
			_maxOddsLogScore = 0.0;
		}

		/// <summary>
		/// Performs the Expectation-Maximization steps. Lets get number-crunchy
		/// </summary>
		/// <param name="startingAlignments">Number of initial starting alignments to use as seeds</param>
		/// <param name="iterations">Number of times to E-M each starting alignment</param>
		public void EM(int numStartingAlignments, int iterations)
		{
			GenerateStartingAlignments(numStartingAlignments);
			int index = 0;

			//For each starting alignment
			foreach (int[] startAlign in _startingAlignments)
			{
				
				//populate the sequence list with the starting alignment
				PopulateStartingAlignment(startAlign);

				//E-M the shit out of it
				for (int i = 0; i < iterations; i++)
				{
					CountResidues(_sequences);
					CalculateFrequencies();
					GenerateLogOdds();
					Maximize();
					CheckMaxLogOddsScore();
				}
				Console.Clear();
				Console.WriteLine("Alignment: {0}/{1}", ++index, numStartingAlignments);
				Console.WriteLine("Maximum Log-Odds Score so far: {0}", _maxOddsLogScore);
				Console.WriteLine("Optimal motifs found so far:\n");
				PrintAlignment(_maxOddsLogScoreSequences); 
			}

			Console.Clear();
			Console.WriteLine("Final Maximum Log-Odds Score: {0}", _maxOddsLogScore);
			Console.WriteLine("Final optimal motifs found:\n");
			Console.WriteLine();
			PrintAlignment(_maxOddsLogScoreSequences);
			
		}

		/// <summary>
		/// Prints the alignment that produced the highest log-odds score all pretty-like
		/// </summary>
		private void PrintAlignment(List<HelperFunctions.Sequence> sequences)
		{

			int index;
			foreach (HelperFunctions.Sequence sequence in sequences)
			{
				index = 0;

				//Print out the sequence up until the motif
				while (index < sequence.motifStart)
					Console.Write(sequence.seq[index++]);

				//Now make the color all fancy for easy readitude
				Console.ForegroundColor = ConsoleColor.Red;
				while (index < (sequence.motifStart + _motifSize))
					Console.Write(sequence.seq[index++]);

				//now reset the colors and continue printing the rest of the sequence
				Console.ResetColor();
				while (index < sequence.seq.Length)
					Console.Write(sequence.seq[index++]);

				Console.WriteLine();
				//And we are done

			}

			//I want to keey the cursor at the top of the screen and have extra sequences bleed off the bottom
			Console.SetCursorPosition(0, 0);

		}

		/// <summary>
		/// Keeps track of the maximum log odds score obtained
		/// </summary>
		private void CheckMaxLogOddsScore()
		{
			//go through all of the sequences comparing log scores
			for (int i = 0; i < _sequences.Count; i++)
			{
				int motifStart = _sequences[i].motifStart;
				double logodd = CalculateLogOdds(_sequences[i].seq.Substring(motifStart, _motifSize));

				if (logodd > _maxOddsLogScore)
				{
					_maxOddsLogScore = logodd;
					_maxOddsLogScoreSequences = _sequences;
				}

			}
		}

		/// <summary>
		/// Populates the sequence list with the starting alignment
		/// </summary>
		/// <param name="startAlignment"></param>
		private void PopulateStartingAlignment(int[] startAlignment)
		{

			HelperFunctions.Sequence tmp;
			for (int i = 0; i < _sequences.Count; i++)
			{
				tmp = _sequences[i];
				tmp.motifStart = startAlignment[i];
				_sequences[i] = tmp;
			}
		}

		/// <summary>
		/// Generates the specified number of starting alignments for the sequences
		/// </summary>
		/// <param name="count"></param>
		private void GenerateStartingAlignments(int count)
		{
			Random rng = new Random();
			List<int[]> startingAlignments = new List<int[]>();
			for (int i = 0; i < count; i++)
			{
				int[] randAlign = new int[_sequences.Count];

				for (int j = 0; j < randAlign.Length; j++)
					//generate a motif start between zero and the length of the sequence - motifSize + 1
					randAlign[j] = rng.Next(_sequences[j].seq.Length - _motifSize + 1);

				//Awwww yeah son add that shit to the list
				startingAlignments.Add(randAlign);
			}

			_startingAlignments = startingAlignments;
		}

		private void CountResidues(List<HelperFunctions.Sequence> sequences)
		{
			//for each string in the list of sequences go through each character and increment that character's count in the dictionary
			foreach (HelperFunctions.Sequence sequence in sequences)
			{

				//initialize the counts
				_residueCounts.Clear();

				for (int i = 0; i < sequence.seq.Length; i++)
				{
					if (_residueCounts.ContainsKey("" + sequence.seq[i]))
						_residueCounts["" + sequence.seq[i]][0]++;
					else
					{
						int[] tmp = new int[_motifSize+1];
						tmp[0] = 1;
						_residueCounts.Add("" + sequence.seq[i], tmp);
					}
				}
			}

			//Now get a count of the residues in the particular positions in the motifs
			foreach (HelperFunctions.Sequence sequence in sequences)
			{
				for (int i = 0; i < _motifSize; i++)
				{
					//Console.Write("" + sequence.seq[sequence.motifStart + i]);
					//add to motif spot count
					_residueCounts[("" + sequence.seq[sequence.motifStart + i])][i+1]++;
					//subtract from background count
					_residueCounts[("" + sequence.seq[sequence.motifStart + i])][0]--;
					//Great job kid! Now lets test this shit and go home.
				}

				//Console.WriteLine();
			}
		
			//Add a pseudocount to all the counts. I am not sure wtf this is but the homework asks for it so whatever
			foreach (string key in _residueCounts.Keys)
				for (int i = 0; i <= _motifSize; i++)
					_residueCounts[key][i]++;

		}

		/// <summary>
		/// Calculates the percentages a given letter appears at a given spot compared to other letters
		/// Can only be called after CountResidues()
		/// </summary>
		private void CalculateFrequencies()
		{

			//Initialize the dictionary
			_residueFrequencies.Clear();

			//initialize the frequency arrays
			//boy the way I data-structured this out sure makes this part just awakward enough to make me feel stoopid
			foreach( string acid in _residueCounts.Keys)
			{
				double[] frequencies = new double[_motifSize + 1];
				_residueFrequencies.Add(acid, frequencies);
			}


			for (int i = 0; i <= _motifSize; i++)
			{
				int total = 0;

				//get the total across all acids
				foreach (string acid in _residueCounts.Keys)
					total += _residueCounts[acid][i];

				//calculate the frequency percentage using the total
				foreach (string acid in _residueCounts.Keys)
				{
					double[] tmp = _residueFrequencies[acid];
					tmp[i] = ((double)_residueCounts[acid][i] / total);
					_residueFrequencies[acid] = tmp;
					//Lets see if this shit works. Unfortunately the test data I'm using has incorrect statistics on it. Ballz
				}
			}
		}

		/// <summary>
		/// From the frequency table take the log base 2 of the frequency of each nucleotide at that position divided by the frequency of it not being in the motif at all
		///	log base 2 of observed divided by expected
		/// </summary>
		private void GenerateLogOdds()
		{
			
			//initialize the dictionary
			_logOddsScores.Clear();

			//	From the frequency table take the log base 2 of the frequency of each nucleotide at that position divided by the frequency of it not being in the motif at all
			//	log base 2 of observed divided by expected

			foreach (string acid in _residueCounts.Keys)
			{
				double[] logOdds = new double[_motifSize + 1];
				for (int i = 1; i <= _motifSize; i++)
				{
					//log base 2 of frequency of position i divided by frequency of position 0 (background, or not in motif)
					logOdds[i] = Math.Log((_residueFrequencies[acid][i] / _residueFrequencies[acid][0]), 2);
					
					//Lets see if this shit works. Unfortunately the test data I'm using has incorrect statistics on it. Ballz
				}
				_logOddsScores.Add(acid, logOdds);
			}

		}

		/// <summary>
		/// Calculates the log odds for a specific motif
		/// </summary>
		private double CalculateLogOdds(string motif)
		{
			double logOdds = 0.0;

			for (int i = 0; i < motif.Length; i++)
			{
				string acid = ("" + motif[i]);
				//add the log odds of that specific acid at that specific position in the motif
				logOdds += _logOddsScores[acid][i + 1]; //i+1 because first cell is not used
			}

			return logOdds;
		}

		public void Maximize()
		{
			double maxLogOdds, currLogOdds;
			int maxMotifStart;
			string motifToTest;
			//calculate the log odds for the current motif
			for(int i = 0; i < _sequences.Count; i++)
			{
				string motif = _sequences[i].seq.Substring(_sequences[i].motifStart, _motifSize);
				maxLogOdds = CalculateLogOdds(motif);
				maxMotifStart = _sequences[i].motifStart;

				//go through every other possible motif calculating the log odds. Length-motifSize+1 because the last bunch characters
				//in the sequence can't be used to start a motif. It would not be long enough
				for (int j = 0; j < _sequences[i].seq.Length - _motifSize + 1; j++)
				{
					motifToTest = _sequences[i].seq.Substring(j, _motifSize);
					currLogOdds = CalculateLogOdds(motifToTest);
					//if we get a better score adjust accordingly
					if (currLogOdds > maxLogOdds)
					{
						maxLogOdds = currLogOdds;
						maxMotifStart = j;
					}
				}

				//set the new motif starting point to the motif with the highest log odds score
				HelperFunctions.Sequence tmp = _sequences[i];
				tmp.motifStart = maxMotifStart;
				_sequences[i] = tmp;

				//???

				//profit

				//fuck bitches, get money
			}
			
		}
	}
}
